#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
Based on Chris Campbell's tutorial from iforce2d.net:
http://www.iforce2d.net/b2dtut/top-down-car
"""

from Box2D.examples.framework import (Framework, Keys, main)
from Box2D import (b2DrawExtended, b2Vec2)
from sim.ground import TDGroundArea
from sim.tire import TDTire
from sim.car import TDCar
from sim.gui import RobotSimulatorGUI

from pygame.locals import (QUIT, KEYDOWN, KEYUP, MOUSEBUTTONDOWN,
                           MOUSEBUTTONUP, MOUSEMOTION, KMOD_LSHIFT)

import math, pygame


try:
    from Box2D.examples.backends.pygame_gui import (fwGUI, gui)
    GUIEnabled = True
except Exception as ex:
    print('Unable to load PGU; menu disabled.')
    print('(%s) %s' % (ex.__class__.__name__, ex))
    GUIEnabled = False


#----------------------------------------------------------------------
class PyGameFramework(Framework):
    TEXTLINE_START = 0

    def __init__(self):
        # 这里要先设置一下屏幕，否则可能会显示偏离
        screen_rec = (1200, 620)
        self.screen = pygame.display.set_mode(screen_rec)
        super(PyGameFramework, self).__init__()

        # Screen and debug draw
        self.screen = pygame.display.set_mode(screen_rec)
        self.screenSize = b2Vec2(*self.screen.get_size())
        self.viewZoom = 3
        self.viewCenter = (110, 100)

        pygame.display.set_caption('Robot Simulator')
        self.settings.drawStats = False
        self.settings.drawFPS = False
        # self.settings.drawMenu = False

        self.font = pygame.font.Font("SimHei.ttf", 24)

        # GUI Initialization
        if GUIEnabled:
            self.gui_app = None
            self.gui_table = None
            self.gui_app = gui.App()
            self.gui_table = RobotSimulatorGUI(self.font)
            container = gui.Container(align = -1, valign = -1)
            container.add(self.gui_table, 0, 0)
            self.gui_app.init(container)


    def checkEvents(self):
        """
        Check for pygame events (mainly keyboard/mouse events).
        Passes the events onto the GUI also.
        """
        for event in pygame.event.get():
            if event.type == QUIT or (event.type == KEYDOWN and event.key == Keys.K_ESCAPE):
                return False
            elif event.type == KEYDOWN:
                self._Keyboard_Event(event.key, down=True)
            elif event.type == KEYUP:
                self._Keyboard_Event(event.key, down=False)
            elif event.type == MOUSEBUTTONDOWN:
                p = self.ConvertScreenToWorld(*event.pos)
                if event.button == 1:  # left
                    mods = pygame.key.get_mods()
                    if mods & KMOD_LSHIFT:
                        self.ShiftMouseDown(p)
                    else:
                        self.MouseDown(p)
                elif event.button == 2:  # middle
                    pass
                elif event.button == 3:  # right
                    self.rMouseDown = True
                elif event.button == 4:
                    # self.viewZoom *= 1.1
                    pass
                elif event.button == 5:
                    # self.viewZoom /= 1.1
                    pass
            elif event.type == MOUSEBUTTONUP:
                p = self.ConvertScreenToWorld(*event.pos)
                if event.button == 3:  # right
                    self.rMouseDown = False
                else:
                    self.MouseUp(p)
            elif event.type == MOUSEMOTION:
                p = self.ConvertScreenToWorld(*event.pos)

                self.MouseMove(p)

                if self.rMouseDown:
                    # self.viewCenter -= (event.rel[0] / 5.0, -event.rel[1] / 5.0)
                    pass

            if GUIEnabled:
                self.gui_app.event(event)  # Pass the event to the GUI

        return True


    def _Keyboard_Event(self, key, down=True):
        """
        Internal keyboard event, don't override this.

        Checks for the initial keydown of the basic testbed keys. Passes the unused
        ones onto the test via the Keyboard() function.
        """
        if down:
            # if key == Keys.K_z:       # Zoom in
            #     self.viewZoom = min(1.1 * self.viewZoom, 50.0)
            # elif key == Keys.K_x:     # Zoom out
            #     self.viewZoom = max(0.9 * self.viewZoom, 0.02)
            # elif key == Keys.K_SPACE:  # Launch a bomb
            #     self.LaunchRandomBomb()
            if key == Keys.K_F1:    # Toggle drawing the menu
                self.settings.drawMenu = not self.settings.drawMenu
            elif key == Keys.K_F2:    # Do a single step
                self.settings.singleStep = True
                if GUIEnabled:
                    self.gui_table.updateGUI(self.settings)
            else:              # Inform the test of the key press
                self.Keyboard(key)
        else:
            self.KeyboardUp(key)

    def CheckKeys(self):
        """
        Check the keys that are evaluated on every main loop iteration.
        I.e., they aren't just evaluated when first pressed down
        """

        pygame.event.pump()
        self.keys = keys = pygame.key.get_pressed()
        # if keys[Keys.K_LEFT]:
        #     self.viewCenter -= (0.5, 0)
        # elif keys[Keys.K_RIGHT]:
        #     self.viewCenter += (0.5, 0)
        #
        # if keys[Keys.K_UP]:
        #     self.viewCenter += (0, 0.5)
        # elif keys[Keys.K_DOWN]:
        #     self.viewCenter -= (0, 0.5)
        #
        # if keys[Keys.K_HOME]:
        #     self.viewZoom = 1.0
        #     self.viewCenter = (0.0, 20.0)


#----------------------------------------------------------------------
class TopDownCar(PyGameFramework):
    name = ""
    description = ""

    def __init__(self):
        super(TopDownCar, self).__init__()
        # Top-down -- no gravity in the screen plane
        self.world.gravity = (0, 0)

        self.key_map = {Keys.K_w: 'up',
                        Keys.K_s: 'down',
                        Keys.K_a: 'left',
                        Keys.K_d: 'right',
                        }

        # Keep track of the pressed keys
        self.pressed_keys = set()

        # The walls
        boundary = self.world.CreateStaticBody(position=(0, 0))
        boundary.CreateEdgeChain([(0, 0),
                                  (300, 0),
                                  (300, 200),
                                  (0, 200),
                                  (0, 0)]
                                 )

        # A couple regions of differing traction
        self.car = TDCar(self.world, position = (25, 25), density = 1.0)
        self.CreateGround()


    def CreateGround(self):
        for idx in range(24):
            y, x = divmod(idx, 6)
            gnd1 = self.world.CreateStaticBody(userData={'obj': TDGroundArea(1)})
            fixture = gnd1.CreatePolygonFixture(box = (25, 25, (25 + 50 * x, 25 + 50 * y), 0))
            fixture.sensor = True


    def Keyboard(self, key):
        key_map = self.key_map
        if key in key_map:
            self.pressed_keys.add(key_map[key])
        else:
            super(TopDownCar, self).Keyboard(key)

    def KeyboardUp(self, key):
        key_map = self.key_map
        if key in key_map:
            self.pressed_keys.remove(key_map[key])
        else:
            super(TopDownCar, self).KeyboardUp(key)

    def handle_contact(self, contact, began):
        # A contact happened -- see if a wheel hit a
        # ground area
        fixture_a = contact.fixtureA
        fixture_b = contact.fixtureB

        body_a, body_b = fixture_a.body, fixture_b.body
        ud_a, ud_b = body_a.userData, body_b.userData
        if not ud_a or not ud_b:
            return

        tire = None
        ground_area = None
        for ud in (ud_a, ud_b):
            obj = ud['obj']
            if isinstance(obj, TDTire):
                tire = obj
            elif isinstance(obj, TDGroundArea):
                ground_area = obj

        if ground_area is not None and tire is not None:
            if began:
                tire.add_ground_area(ground_area)
            else:
                tire.remove_ground_area(ground_area)

    def BeginContact(self, contact):
        self.handle_contact(contact, True)

    def EndContact(self, contact):
        self.handle_contact(contact, False)

    def ShiftMouseDown(self, p):
        """
        Indicates that there was a left click at point p (world coordinates)
        with the left shift key being held down.
        """
        self.mouseWorld = p

        if not self.mouseJoint:
            self.SpawnBomb(p)

    def MouseDown(self, p):
        """
        Indicates that there was a left click at point p (world coordinates)
        """
        if self.mouseJoint is not None:
            return

        # Create a mouse joint on the selected body (assuming it's dynamic)
        # Make a small box.
        # aabb = b2AABB(lowerBound=p - (0.001, 0.001),
        #               upperBound=p + (0.001, 0.001))
        #
        # # Query the world for overlapping shapes.
        # query = fwQueryCallback(p)
        # self.world.QueryAABB(query, aabb)
        #
        # if query.fixture:
        #     body = query.fixture.body
        #     # A body was selected, create the mouse joint
        #     self.mouseJoint = self.world.CreateMouseJoint(
        #         bodyA=self.groundbody,
        #         bodyB=body,
        #         target=p,
        #         maxForce=1000.0 * body.mass)
        #     body.awake = True

    def MouseUp(self, p):
        """
        Left mouse button up.
        """
        if self.mouseJoint:
            self.world.DestroyJoint(self.mouseJoint)
            self.mouseJoint = None

        if self.bombSpawning:
            self.CompleteBombSpawn(p)

    def MouseMove(self, p):
        """
        Mouse moved to point p, in world coordinates.
        """
        self.mouseWorld = p
        if self.mouseJoint:
            self.mouseJoint.target = p

    def Step(self, settings):
        self.car.update(self.pressed_keys, settings.hz)

        super(TopDownCar, self).Step(settings)

        tractions = [tire.current_traction for tire in self.car.tires]
        # self.Print('Current tractions: %s' % tractions)

if __name__ == "__main__":
    main(TopDownCar)
